from typing import Generator

from openai import OpenAI


class VllmModel:

    def __init__(self, api_key="EMPTY", base_url="http://172.16.3.116:2380/v1",
                 model="Qwen3-32B-AWQ", enable_thinking=None):
        self.client = OpenAI(api_key=api_key, base_url=base_url)
        self.model = model
        self.enable_thinking = enable_thinking

    def generate(self, messages, temperature=0.6) -> str:
        if self.enable_thinking is not None:
            res = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                temperature=temperature,
                stream=False,
                extra_body={"chat_template_kwargs": {"enable_thinking": self.enable_thinking}}
            )
            return res.choices[0].message.content
        res = self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            temperature=temperature,
            stream=False
        )

        return res.choices[0].message.content

    def stream(self, messages, temperature=0.3) -> Generator:

        res = self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            temperature=temperature,
            stream=True,
            extra_body={"chat_template_kwargs": {"enable_thinking": self.enable_thinking}}
        )
        for chunk in res:
            content = chunk.choices[0].delta.content
            if content is not None:
                yield content
